package com.example.aichat.service;

import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;

import javax.annotation.PostConstruct;
import java.util.concurrent.atomic.AtomicBoolean;

@Service
public class OllamaService {

    @Value("${ai.ollama.base_url}")
    private String ollama_base_url;

    @Value("${ai.ollama.model}")
    private String ollama_model;

    private WebClient webClient;

    @PostConstruct
    public void init() {
        this.webClient = WebClient.builder()
                .baseUrl(ollama_base_url)
                .build();
    }

    public void streamChat(String content, ResponseBodyEmitter emitter) {

        Flux<String> responseFlux = webClient.post()
                .uri("/api/generate")
                .bodyValue(new OllamaRequest(ollama_model, content))
                .retrieve()
                .bodyToFlux(OllamaResponse.class)
                .map(response -> response.response)
                .map(response -> {
//                    System.out.println("Response received: " + response);// 处理特殊字符：将换行符、回车符、制表符等转换为它们的转义表示
                    String escapedResponse = response
                            .replace("\n", "\\n")
                            .replace("\r", "\\r")
                            .replace("\t", "\\t")
                            .replace("\"", "\\\"");
                    return String.format("{\"response\": \"%s\"}\n", escapedResponse);
                })
                ;

        responseFlux.subscribe(
                data -> {
                    try {
//                        System.out.println("Sending data: " + data);
                        emitter.send(data, MediaType.APPLICATION_JSON);
                    } catch (Exception e) {
                        emitter.completeWithError(e);
                    }
                },
                emitter::completeWithError,
                emitter::complete);
    }


    private record OllamaRequest(String model, String prompt) {}
    private record OllamaResponse(String response) {}
}